// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <spica/compressor.h>
#include <spica/config.h>
#include <zlib.h>

namespace spica::detail {

    class gzip_compressor final : public compressor {
    public:
        gzip_compressor();
        ~gzip_compressor() override;

        bool compress(const char *data, size_t data_length, bool last,
                      Callback callback) override;

    private:
        bool is_valid_ = false;
        z_stream strm_;
    };

    class gzip_decompressor final : public decompressor {
    public:
        gzip_decompressor();
        ~gzip_decompressor() override;

        bool is_valid() const override;

        bool decompress(const char *data, size_t data_length,
                        Callback callback) override;

    private:
        bool is_valid_ = false;
        z_stream strm_;
    };

    inline gzip_compressor::gzip_compressor() {
        std::memset(&strm_, 0, sizeof(strm_));
        strm_.zalloc = Z_NULL;
        strm_.zfree = Z_NULL;
        strm_.opaque = Z_NULL;

        is_valid_ = deflateInit2(&strm_, Z_DEFAULT_COMPRESSION, Z_DEFLATED, 31, 8,
                                 Z_DEFAULT_STRATEGY) == Z_OK;
    }

    inline gzip_compressor::~gzip_compressor() { deflateEnd(&strm_); }

    inline bool gzip_compressor::compress(const char *data, size_t data_length,
                                          bool last, Callback callback) {
        assert(is_valid_);

        do {
            constexpr size_t max_avail_in =
                    (std::numeric_limits<decltype(strm_.avail_in)>::max)();

            strm_.avail_in = static_cast<decltype(strm_.avail_in)>(
                    (std::min)(data_length, max_avail_in));
            strm_.next_in = const_cast<Bytef *>(reinterpret_cast<const Bytef *>(data));

            data_length -= strm_.avail_in;
            data += strm_.avail_in;

            auto flush = (last && data_length == 0) ? Z_FINISH : Z_NO_FLUSH;
            auto ret = Z_OK;

            std::array<char, SPICA_COMPRESSION_BUFSIZ> buff{};
            do {
                strm_.avail_out = static_cast<uInt>(buff.size());
                strm_.next_out = reinterpret_cast<Bytef *>(buff.data());

                ret = deflate(&strm_, flush);
                if (ret == Z_STREAM_ERROR) { return false; }

                if (!callback(buff.data(), buff.size() - strm_.avail_out)) {
                    return false;
                }
            } while (strm_.avail_out == 0);

            assert((flush == Z_FINISH && ret == Z_STREAM_END) ||
                   (flush == Z_NO_FLUSH && ret == Z_OK));
            assert(strm_.avail_in == 0);
        } while (data_length > 0);

        return true;
    }

    inline gzip_decompressor::gzip_decompressor() {
        std::memset(&strm_, 0, sizeof(strm_));
        strm_.zalloc = Z_NULL;
        strm_.zfree = Z_NULL;
        strm_.opaque = Z_NULL;

        // 15 is the value of wbits, which should be at the maximum possible value
        // to ensure that any gzip stream can be decoded. The offset of 32 specifies
        // that the stream type should be automatically detected either gzip or
        // deflate.
        is_valid_ = inflateInit2(&strm_, 32 + 15) == Z_OK;
    }

    inline gzip_decompressor::~gzip_decompressor() { inflateEnd(&strm_); }

    inline bool gzip_decompressor::is_valid() const { return is_valid_; }

    inline bool gzip_decompressor::decompress(const char *data, size_t data_length,
                                              Callback callback) {
        assert(is_valid_);

        auto ret = Z_OK;

        do {
            constexpr size_t max_avail_in =
                    (std::numeric_limits<decltype(strm_.avail_in)>::max)();

            strm_.avail_in = static_cast<decltype(strm_.avail_in)>(
                    (std::min)(data_length, max_avail_in));
            strm_.next_in = const_cast<Bytef *>(reinterpret_cast<const Bytef *>(data));

            data_length -= strm_.avail_in;
            data += strm_.avail_in;

            std::array<char, SPICA_COMPRESSION_BUFSIZ> buff{};
            while (strm_.avail_in > 0 && ret == Z_OK) {
                strm_.avail_out = static_cast<uInt>(buff.size());
                strm_.next_out = reinterpret_cast<Bytef *>(buff.data());

                ret = inflate(&strm_, Z_NO_FLUSH);

                assert(ret != Z_STREAM_ERROR);
                switch (ret) {
                    case Z_NEED_DICT:
                    case Z_DATA_ERROR:
                    case Z_MEM_ERROR: inflateEnd(&strm_); return false;
                }

                if (!callback(buff.data(), buff.size() - strm_.avail_out)) {
                    return false;
                }
            }

            if (ret != Z_OK && ret != Z_STREAM_END) { return false; }

        } while (data_length > 0);

        return true;
    }

}  // namespace spica::detail

